﻿using System;
using System.Collections.Generic;
using System.IO;

namespace ooWii.IO
{
	public class WbfsDiscList : List<WbfsDiscInfo>
	{
		#region WbfsDiscList

		const int wbfs_disc_info_size = 0x100;
		const UInt32 n_wii_sec_per_disc = 143432 * 2;
		const int wii_sec_sz = 0x8000;
		const int wii_sec_sz_s = 15;

		public WbfsDiscList( Stream stream )
		{
			byte[] wbfsheader_buffer = new byte[ 12 ];

			stream.Position = 0;
			if( stream.Read( wbfsheader_buffer, 0, wbfsheader_buffer.Length ) != wbfsheader_buffer.Length )
			{
				throw new WiiException( "Cannot read wbfs header" );
			}

			if( WiiBuffer.GetString( wbfsheader_buffer, 0, 4 ) != "WBFS" )
			{
				throw new Exception( "not a wbfs partition" );
			}

			UInt32 n_hd_sec = WiiBuffer.GetUInt32( wbfsheader_buffer, 4 );
			byte hd_sec_sz_s = wbfsheader_buffer[ 8 ];
			byte wbfs_sec_sz_s = wbfsheader_buffer[ 9 ];

			UInt32 hd_sec_sz = 1U << hd_sec_sz_s;

			UInt32 n_wii_sec = ( n_hd_sec / wii_sec_sz ) * hd_sec_sz;
			UInt32 wbfs_sec_sz = 1U << wbfs_sec_sz_s;
			UInt32 n_wbfs_sec = n_wii_sec >> ( wbfs_sec_sz_s - wii_sec_sz_s );
			UInt32 n_wbfs_sec_per_disc = n_wii_sec_per_disc >> ( wbfs_sec_sz_s - wii_sec_sz_s );
			int disc_info_sz = (int)sectorAlign( wbfs_disc_info_size + ( n_wbfs_sec_per_disc * 2 ), hd_sec_sz );

			UInt32 fb_memsize = sectorAlign( n_wbfs_sec / 8, hd_sec_sz );
			UInt32 freeblks_lba = (UInt32)( ( wbfs_sec_sz - ( n_wbfs_sec / 8 ) ) >> hd_sec_sz_s );
			//UInt32 freeblks_sectors = fb_memsize >> hd_sec_sz_s;

			n_wii_sec = n_hd_sec / ( wii_sec_sz / hd_sec_sz );
			n_wbfs_sec = n_wii_sec >> ( wbfs_sec_sz_s - wii_sec_sz_s );
			if( n_wbfs_sec > 0x10000 )
			{
				n_wbfs_sec = 0x10000;
			}

			UInt32 freeblks_size4 = ( ( n_wbfs_sec - 1 ) + 31 ) / 32;
			if( freeblks_size4 > fb_memsize / 4 )
			{
				freeblks_size4 = fb_memsize / 4;
				UInt32 max_sec = ( freeblks_size4 * 32 ) + 1;
				if( n_wbfs_sec > max_sec )
				{
					n_wbfs_sec = max_sec;
				}
			}

			//UInt32 freeblks_mask = (1U << (int)((n_wbfs_sec - 1) & 31)) - 1;

			UInt32 max_disc = (UInt32)( ( freeblks_lba - 1 ) / ( disc_info_sz >> hd_sec_sz_s ) );

			UInt32 max_slots = (UInt32)( hd_sec_sz - wbfsheader_buffer.Length );
			if( max_disc > max_slots )
			{
				max_disc = max_slots;
			}

			byte[] slotused_buffer = new byte[ max_disc ];
			stream.Position = wbfsheader_buffer.Length;
			if( stream.Read( slotused_buffer, 0, slotused_buffer.Length ) != slotused_buffer.Length )
			{
				throw new WiiException( "Cannot read used slot table" );
			}

			byte[] freeBlocksBuffer = new byte[ fb_memsize ];
			stream.Position = freeblks_lba * hd_sec_sz;
			if( stream.Read( freeBlocksBuffer, 0, freeBlocksBuffer.Length ) != fb_memsize )
			{
				throw new WiiException( "Cannot read free blocks bitmap" );
			}

			UInt32[] freeBlocks = new UInt32[ fb_memsize / 4 ];
			for( int i = 0; i < freeBlocks.Length; i++ )
			{
				freeBlocks[ i ] = WiiBuffer.GetUInt32( freeBlocksBuffer, i * 4 );
			}

			UInt32[] calculatedFreeBlocks = new UInt32[ fb_memsize / 4 ];
			for( int i = 1; i < n_wbfs_sec; i++ )
			{
				calculatedFreeBlocks[ ( i - 1 ) / 32 ] |= 1U << ( ( i - 1 ) % 32 );
			}

			for( int disc = 0; disc < max_disc; disc++ )
			{
				if( slotused_buffer[ disc ] != 0 )
				{
					byte[] discinfo_buffer = new byte[ wbfs_disc_info_size ];

					stream.Position = hd_sec_sz + ( disc * disc_info_sz );
					if( stream.Read( discinfo_buffer, 0, discinfo_buffer.Length ) != discinfo_buffer.Length )
					{
						throw new WiiException( "Cannot read disc info" );
					}

					if( WiiBuffer.GetUInt32( discinfo_buffer, 24 ) == 0x5d1c9ea3 )
					{
						byte[] blockMapBuffer = new byte[ n_wbfs_sec_per_disc * 2 ];
						if( stream.Read( blockMapBuffer, 0, blockMapBuffer.Length ) != blockMapBuffer.Length )
						{
							throw new WiiException( "Cannot read block map" );
						}

						for( int logicalBlock = 0; logicalBlock < blockMapBuffer.Length / 2; logicalBlock++ )
						{
							UInt16 physicalBlock = WiiBuffer.GetUInt16( blockMapBuffer, logicalBlock * 2 );
							if( physicalBlock != 0 )
							{
								if( physicalBlock >= n_wbfs_sec )
								{
									throw new WiiException( string.Format( "block {0} out of range", physicalBlock ) );
								}

								int freeBlockIndex = ( physicalBlock - 1 ) / 32;
								UInt32 freeBlockBit = 1U << ( ( physicalBlock - 1 ) % 32 );

								if( ( calculatedFreeBlocks[ freeBlockIndex ] & freeBlockBit ) == 0 )
								{
									throw new WiiException( string.Format( "block {0} reused", physicalBlock ) );
								}

								calculatedFreeBlocks[ freeBlockIndex ] &= ~freeBlockBit;
							}
						}

						this.Add( new WbfsDiscInfo( stream, discinfo_buffer, blockMapBuffer, wbfs_sec_sz ) );
					}
				}
			}

			for( int i = 0; i < freeBlocks.Length; i++ )
			{
				if( freeBlocks[ i ] != calculatedFreeBlocks[ i ] )
				{
					throw new WiiException( "free block bitmap incorrect" );
				}
			}
		}

		static UInt32 sectorAlign( UInt32 offset, UInt32 blockSize )
		{
			blockSize--;
			return ( offset + blockSize ) & ~blockSize;
		}

		#endregion
	}
}
